use crate::Tensor;

pub struct FunctionCtx<T> {
    forward_inputs: Vec<Tensor<T>>,
    tensors: Option<Vec<Tensor<T>>>,
    others: Option<Vec<Vec<usize>>>,
}

impl<T> FunctionCtx<T> {
    pub fn new(forward_inputs: Vec<Tensor<T>>) -> FunctionCtx<T> {
        FunctionCtx {
            forward_inputs,
            tensors: None,
            others: None,
        }
    }

    pub fn forward_inputs(&self) -> &Vec<Tensor<T>> {
        &self.forward_inputs
    }

    pub fn others(&self) -> Vec<&[usize]> {
        self.others
            .as_ref()
            .map(|x| x.iter().map(|x| x.as_slice()))
            .unwrap()
            .collect()
    }

    pub fn tensors(&self) -> &Vec<Tensor<T>> {
        self.tensors.as_ref().unwrap()
    }

    pub fn save_tensors(&mut self, tensors: Vec<Tensor<T>>) {
        self.tensors = Some(tensors);
    }

    pub fn save_others(&mut self, foo: Vec<&[usize]>) {
        self.others = foo
            .iter()
            .map(|x| x.to_vec())
            .collect::<Vec<Vec<usize>>>()
            .into();
    }
}
